/**
 * @Author yang
 * @Description 协程池工具类, 用于量化的管理协程创建, 控制系统资源的使用
 * @Date 2021/1/8
 * @Time 9:31
 **/
package pool

import (
	"encoding/json"
	"errors"
	"log"
	"sync"
	"sync/atomic"
	"time"
)

var (
	// ErrInvalidPoolCap return if pool size <= 0
	ErrInvalidPoolCap = errors.New("invalid pool cap")
	// ErrInvalidPoolRoutineSize return if pool routine size <= 0
	ErrInvalidPoolRoutineSize = errors.New("invalid pool routine size")
	// ErrPoolAlreadyClosed put task but pool already closed
	ErrPoolAlreadyClosed = errors.New("pool already closed")
)

const (
	// RUNNING pool is running
	RUNNING = 1
	// STOPED pool is stoped
	STOPED = 0
)

// Task task to-do
type Task struct {
	TaskName string                 // 任务名称
	Handler  func(v ...interface{}) // 任务执行函数
	Params   []interface{}          // 执行函数参数
}

func NewTask(taskname string, handler func(v ...interface{}), params []interface{}) *Task {
	return &Task{
		TaskName: taskname,
		Handler:  handler,
		Params:   params,
	}
}

// Pool task pool
type Pool struct {
	routineSize    uint64            // 协程数
	capacity       uint64            // 最大任务容量
	runningWorkers uint64            // 实际运行中的协程数
	state          int64             // 协程池状态， 1运行|0停止
	taskC          chan *Task        // taskChannel
	PanicHandler   func(interface{}) // 异常处理函数
	sync.Mutex
}

// NewPool init pool
func NewPool(capacity uint64, routineSize uint64) (*Pool, error) {
	if capacity <= 0 {
		return nil, ErrInvalidPoolCap
	}
	if capacity <= 0 {
		return nil, ErrInvalidPoolRoutineSize
	}
	return &Pool{
		routineSize: routineSize,
		capacity:    capacity,
		state:       RUNNING,
		taskC:       make(chan *Task, capacity),
	}, nil
}

// GetCap get capacity
func (p *Pool) GetRoutineSize() uint64 {
	return p.routineSize
}

// GetCap get capacity
func (p *Pool) GetCap() uint64 {
	return p.capacity
}

// GetRunningWorkers get running workers
func (p *Pool) GetRunningWorkers() uint64 {
	return atomic.LoadUint64(&p.runningWorkers)
}

func (p *Pool) incRunning() {
	atomic.AddUint64(&p.runningWorkers, 1)
}

func (p *Pool) decRunning() {
	atomic.AddUint64(&p.runningWorkers, ^uint64(0))
}

// Put put a task to pool
func (p *Pool) Put(task *Task) (err error) {

	if p.getState() == STOPED {
		return ErrPoolAlreadyClosed
	}

	// safe run worker
	p.Lock()
	if p.GetRunningWorkers() < p.GetRoutineSize() {
		p.run()
	}
	p.Unlock()

	// send task safe
	p.Lock()
	if p.state == RUNNING {
		if uint64(len(p.taskC)) < p.GetCap() {
			p.taskC <- task
		} else {
			bytes, _ := json.Marshal(task)
			err = errors.New("===>> 协程池任务已满! 当前要加入的任务信息为:" + string(bytes))
		}
	}
	p.Unlock()

	return
}

func (p *Pool) run() {
	p.incRunning()

	go func() {
		defer func() {
			p.decRunning()
			if r := recover(); r != nil {
				if p.PanicHandler != nil {
					p.PanicHandler(r)
				} else {
					log.Printf("Worker panic: %s\n", r)
				}
			}
		}()

		for {
			select {
			case task, ok := <-p.taskC:
				if !ok {
					return
				}
				task.Handler(task.Params...)
			}
		}
	}()
}

func (p *Pool) getState() int64 {
	p.Lock()
	defer p.Unlock()

	return p.state
}

func (p *Pool) setState(state int64) {
	p.Lock()
	defer p.Unlock()

	p.state = state
}

// close safe
func (p *Pool) close() {
	p.Lock()
	defer p.Unlock()

	close(p.taskC)
}

// Close close pool graceful
func (p *Pool) Close() {

	if p.getState() == STOPED {
		return
	}

	p.setState(STOPED) // stop put task

	for len(p.taskC) > 0 { // wait all task be consumed
		time.Sleep(1e6) // reduce CPU load
	}

	p.close()
}
